// ChanceTree.cs
//
// Copyright (c) 2007 Mike Strobel
//
// This source code is subject to the terms of the Microsoft Reciprocal License (Ms-RL).
// For details, see <http://www.opensource.org/licenses/ms-rl.html>.
//
// All other rights reserved.

using System;
using System.Collections.Generic;

namespace Supremacy.Utility
{
    public class Chance<TItem>
    {
        public int GrowthWeight { get; set; }
        public int CurrentWeight { get; set; }
        public TItem Item { get; set; }

        public Chance(int growthWeight, int currentWeight, TItem item)
        {
            this.GrowthWeight = growthWeight;
            this.CurrentWeight = currentWeight;
            this.Item = item;
        }
    }

    public class ChanceTree<TItem>
    {
        private readonly Random _random;
        private readonly List<Chance<TItem>> _chances;

        public ChanceTree()
        {
            _chances = new List<Chance<TItem>>();
            _random = new MersenneTwister();
        }

        protected IList<Chance<TItem>> Chances
        {
            get { return _chances; }
        }

        public void AddChance(int chance, TItem item)
        {
            int index = _chances.Count;
            _chances.Insert(index, new Chance<TItem>(chance, chance, item));
            while (index > 0)
            {
                index = (index - 1) >> 1;
                _chances[index].GrowthWeight += chance;
            }
        }

        public TItem Get()
        {
            if (_chances.Count == 0)
                return default(TItem);
            return _chances[FindNode(_random.Next(_chances[0].GrowthWeight))].Item;
        }

        public TItem Take()
        {
            if (_chances.Count == 0)
                return default(TItem);
            int probability = _random.Next(_chances[0].GrowthWeight);
            int index = FindNode(probability);
            Chance<TItem> chance = _chances[index];
            ModifyChances(index, -chance.CurrentWeight);
            return chance.Item;
        }

        protected void ModifyChances(int index, int modifyBy)
        {
            // Get the prior growth
            int count = _chances.Count;
            int priorGrowth = _chances[index].GrowthWeight;
            _chances[index].CurrentWeight += modifyBy;
            _chances[index].GrowthWeight += modifyBy;
            if (_chances[index].CurrentWeight <= 0)
            {
                count--;
                if (count != index)
                {
                    // First remove the last node from the list
                    int swapIndex = count;
                    while (swapIndex > 0)
                    {
                        swapIndex = (swapIndex - 1) >> 1;
                        _chances[swapIndex].GrowthWeight -= _chances[count].CurrentWeight;
                    }
                    priorGrowth = _chances[index].GrowthWeight;
                    // Swap the last with the current
                    _chances[index] = _chances[count];
                    // Get the new growth weight
                    int son = (index << 1) + 1;
                    _chances[index].GrowthWeight = _chances[index].CurrentWeight;
                    if (son < count)
                    {
                        _chances[index].GrowthWeight += _chances[son++].GrowthWeight;
                    }
                    if (son < count)
                    {
                        _chances[index].GrowthWeight += _chances[son].GrowthWeight;
                    }
                }
            }
            // Feed back up the tree
            int feedWeight = _chances[index].GrowthWeight - priorGrowth;
            if (feedWeight != 0)
            {
                // Parent Weight
                int pIndex = index;
                while (pIndex > 0)
                {
                    pIndex = (pIndex - 1) >> 1;
                    _chances[pIndex].GrowthWeight += feedWeight;
                }
            }
        }

        protected int FindNode(int chance)
        {
            int index = 0;
            int prior = 0;
            while (true)
            {
                prior += _chances[index].CurrentWeight;
                if (chance < prior)
                {
                    return index;
                }
                index = (index << 1) + 1;
                if (chance >= (prior + _chances[index].GrowthWeight))
                {
                    prior += _chances[index].GrowthWeight;
                    index++;
                }
            }
        }
    }
}